package gossip.v2.cluster;

import com.github.benmanes.caffeine.cache.Caffeine;
import gossip.v2.GossipCore;
import gossip.v2.message.IndirectCheckRequest;
import lombok.Getter;
import lombok.Setter;
import org.ricks.ioc.utils.Logger;
import org.ricks.net.transport.AioSession;
import java.util.*;
import java.util.concurrent.*;
// 集群管理器
import java.util.stream.Collectors;

import com.github.benmanes.caffeine.cache.LoadingCache;

@Getter
public class ClusterManager {
    private final GossipNode localNode;
    private final List<GossipNode> seedNodes = new CopyOnWriteArrayList<>();
    private final Map<String, GossipNode> liveNodes = new ConcurrentHashMap<>();
    private final Map<String, GossipNode> suspiciousNodes = new ConcurrentHashMap<>();
    private final Map<AioSession, GossipNode> sessionToNodeMap = new ConcurrentHashMap<>();
    private final Map<String, AioSession> nodeSessionMap = new ConcurrentHashMap<>();
    private final List<NodeStateListener> listeners = new CopyOnWriteArrayList<>();
    private final LoadingCache<String, GossipNode> nodeCache;

    @Setter
    private GossipCore core;
    @Setter private int suspicionTimeout = 5000;  // 改为可配置
    @Setter private int removalTimeout = 10000;

    public ClusterManager(GossipNode localNode) {
        this.localNode = localNode;
        this.liveNodes.put(localNode.getNodeId(), localNode);
        this.nodeCache = Caffeine.newBuilder()
                .expireAfterAccess(30, TimeUnit.SECONDS)
                .maximumSize(1000)
                .build(this::findNodeWithoutCache);
    }

    public GossipNode resolveNode(AioSession session) {
        return sessionToNodeMap.get(session);
    }

    public void registerSession(GossipNode node, AioSession session) {
        sessionToNodeMap.put(session, node);
        nodeSessionMap.put(node.getNodeId(), session);
    }

    public void onHeartbeat(GossipNode node, long timestamp) {
        node.updateHeartbeat(timestamp);
        Logger.debug("{} - 存活得节点：{}  - 收到 {} 发来得心跳", core.getLocalNode().getNodeId(),
                Arrays.toString(liveNodes.keySet().toArray()), node.getNodeId());
        if (!liveNodes.containsKey(node.getNodeId())) {
            moveToLiveNodes(node.getNodeId(), node);
            notifyNodeJoined(node);
        }
    }

    public void addSeedNode(GossipNode node) {
        if (seedNodes.stream().noneMatch(n -> n.getNodeId().equals(node.getNodeId()))) {
            seedNodes.add(node);
            if (!isKnownNode(node.getNodeId())) {
                liveNodes.put(node.getNodeId(), node);
                notifyNodeJoined(node);
            }
            invalidateNodeCache(node.getNodeId());
        }
    }

    public boolean shouldReconnect(GossipNode node) {
        return isSeedNode(node.getNodeId()) ||
                liveNodes.containsKey(node.getNodeId()) ||
                suspiciousNodes.containsKey(node.getNodeId());
    }

    // 添加方法：获取可疑节点
    public Map<String, GossipNode> getSuspiciousNodes() {
        return new ConcurrentHashMap<>(suspiciousNodes);
    }

    public boolean isSeedNode(String nodeId) {
        return seedNodes.stream().anyMatch(n -> n.getNodeId().equals(nodeId));
    }

    public void handleIndirectResponse(String suspiciousNodeId, boolean isAlive) {
        GossipNode node = getNodeById(suspiciousNodeId);
        if (node != null) {
            if (isAlive) {
                moveToLiveNodes(suspiciousNodeId, node);
            } else {
                removeNode(node);
            }
            invalidateNodeCache(suspiciousNodeId);
        }
    }

    public GossipNode getNodeById(String nodeId) {
        try {
            return nodeCache.get(nodeId);
        } catch (Exception e) {
            return findNodeWithoutCache(nodeId);
        }
    }

    // 优化方法：添加发现的节点
    public void addDiscoveredNode(GossipNode node) {
        if (!isKnownNode(node.getNodeId())) {
            Logger.info("Adding discovered node: {}", node.getNodeId());
            liveNodes.put(node.getNodeId(), node);
            notifyNodeJoined(node);
            // 触发主动连接
            core.getTransport().connectToNode(node);
        } else {
            // 更新节点地址
            updateNodeAddress(node.getNodeId(), node.getIp(), node.getPort());
        }
    }

    public Map<String, GossipNode> getAllKnownNodesWithAddress() {
        Map<String, GossipNode> allNodes = new HashMap<>();

        // 添加本地节点
        allNodes.put(localNode.getNodeId(), localNode);

        // 添加活跃节点
        liveNodes.values().forEach(node -> allNodes.put(node.getNodeId(), node));

        // 添加可疑节点
        suspiciousNodes.values().forEach(node -> allNodes.put(node.getNodeId(), node));

        // 添加种子节点
        seedNodes.forEach(node -> {
            if (!allNodes.containsKey(node.getNodeId())) {
                allNodes.put(node.getNodeId(), node);
            }
        });

        return allNodes;
    }

    public boolean isKnownNode(String nodeId) {
        return liveNodes.containsKey(nodeId) ||
                suspiciousNodes.containsKey(nodeId) ||
                isSeedNode(nodeId);
    }

    private GossipNode findNodeWithoutCache(String nodeId) {
        if (localNode.getNodeId().equals(nodeId)) return localNode;
        GossipNode node = liveNodes.get(nodeId);
        if (node != null) return node;
        node = suspiciousNodes.get(nodeId);
        if (node != null) return node;
        return seedNodes.parallelStream()
                .filter(seed -> seed.getNodeId().equals(nodeId))
                .findFirst()
                .orElse(null);
    }

    private void moveToLiveNodes(String nodeId, GossipNode node) {
        suspiciousNodes.remove(nodeId);
        liveNodes.put(nodeId, node);
        node.updateLastHeartbeat();
    }

    public void updateNodeAddress(String nodeId, String ip, int port) {
        GossipNode node = getNodeById(nodeId);
        if (node != null) {
            node.setIp(ip);
            node.setPort(port);
        }
    }

    public void detectFailedNodes() {
        long now = System.currentTimeMillis();

        // 检查活跃节点
        liveNodes.values().stream()
                .filter(n -> !n.isLocal())
                .filter(n -> now - n.getLastHeartbeat() > suspicionTimeout)
                .forEach(this::markAsSuspicious);

        // 检查可疑节点
        suspiciousNodes.values().removeIf(node -> {
            long inactiveTime = now - node.getLastHeartbeat();
            if (inactiveTime > removalTimeout) {
                removeNode(node);
                return true;
            }
            return false;
        });
    }

    public List<GossipNode> selectRandomNodes(int count) {
        List<GossipNode> candidates = liveNodes.values().stream()
                .filter(n -> !n.equals(localNode))
                .filter(n -> !core.getTransport().hasActiveSession(n))
                .collect(Collectors.toList());

        if (candidates.isEmpty()) return candidates;

        // 拓扑感知选择
        candidates.sort(Comparator.comparingDouble(node ->
                core.getTopologyManager().getNodeLatency(node.getNodeId())));

        return candidates.subList(0, Math.min(count, candidates.size()));
    }

    // 新增方法：记录可疑节点
    public void markAsSuspicious(GossipNode node) {
        String nodeId = node.getNodeId();
        if (liveNodes.remove(nodeId) != null) {
            Logger.info("Node:{} - marked as suspicious: {}",
                    core.getLocalNode().getNodeId(), node);
            suspiciousNodes.put(nodeId, node);
            requestIndirectCheck(node);
            invalidateNodeCache(nodeId);
        }
    }

    private void requestIndirectCheck(GossipNode node) {
        String nodeId = node.getNodeId();
        List<GossipNode> proxies = liveNodes.values().stream()
                .filter(n -> !n.getNodeId().equals(nodeId) && !n.isLocal())
                .toList();

        if (!proxies.isEmpty()) {
            GossipNode proxy = proxies.get(ThreadLocalRandom.current().nextInt(proxies.size()));
            core.getTransport().send(proxy, new IndirectCheckRequest(localNode, node));
        }
    }

    public void removeNode(GossipNode node) {
        if (node == null) return;

        String nodeId = node.getNodeId();
        if (liveNodes.remove(nodeId) != null || suspiciousNodes.remove(nodeId) != null) {
            clearNodeSessions(nodeId);
            invalidateNodeCache(nodeId);

            if (!nodeId.equals(localNode.getNodeId())) {
                notifyNodeLeft(node);
            }
        }
    }

    public void initSeedConnections() {
        seedNodes.stream()
                .filter(node -> !node.equals(localNode))
                .filter(node -> core.getTransport().shouldInitiateConnection(localNode, node))
                .forEach(node -> core.getTransport().scheduleReconnect(node));
    }

    private void clearNodeSessions(String nodeId) {
        sessionToNodeMap.entrySet().removeIf(entry -> entry.getValue().getNodeId().equals(nodeId));
        nodeSessionMap.remove(nodeId);
    }

    public void clearNodeSession(AioSession session) {
        GossipNode node = sessionToNodeMap.remove(session);
        if (node != null) {
            nodeSessionMap.remove(node.getNodeId());
        }
    }

    private void invalidateNodeCache(String nodeId) {
        nodeCache.invalidate(nodeId);
    }

    private void notifyNodeJoined(GossipNode node) {
        Logger.info("Node:{} - JOINED: {}", core.getLocalNode().getNodeId(), node);
        listeners.forEach(l -> l.onNodeJoined(node));
    }

    private void notifyNodeLeft(GossipNode node) {
        Logger.info("Node:{} - left: {}", core.getLocalNode().getNodeId(), node);
        listeners.forEach(l -> l.onNodeLeft(node));
    }


    public Set<String> getAllKnownNodes() {
        Set<String> allNodes = new HashSet<>(liveNodes.keySet());
        allNodes.addAll(suspiciousNodes.keySet());
        allNodes.add(localNode.getNodeId());
        seedNodes.stream().map(GossipNode::getNodeId).forEach(allNodes::add);
        return allNodes;
    }

    public void addListener(NodeStateListener listener) {
        listeners.add(listener);
    }

    public List<GossipNode> getSeedNodes() {
        return new ArrayList<>(seedNodes);
    }

    public Map<String, GossipNode> getLiveNodes() {
        return new ConcurrentHashMap<>(liveNodes);
    }

    public interface NodeStateListener {
        void onNodeJoined(GossipNode node);
        void onNodeLeft(GossipNode node);
    }
}
